/*
*    l3qd - Light, light, lightweight queue
*    Copyright (C) 2024  Marcus Pedersén marcus@marcux.org
*
*    This program is free software: you can redistribute it and/or modify
*    it under the terms of the GNU General Public License as published by
*    the Free Software Foundation, either version 3 of the License, or
*    (at your option) any later version.
*
*    This program is distributed in the hope that it will be useful,
*    but WITHOUT ANY WARRANTY; without even the implied warranty of
*    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
*    GNU General Public License for more details.
*
*    You should have received a copy of the GNU General Public License
*    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package node

import (
	"fmt"
	"notabug.org/marcux/l3q/internal/common"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strconv"
	"strings"
)

// Returns current status
// on all running tasks
// If output string is empty
// string, stdout and stderr
// will not be collected.
func AllTaskStatus(sliceName, outputDir string) ([]common.NodeTask, error) {
	tasks := []common.NodeTask{}
	errs := ""

	cmdSystemctl := exec.Command("/usr/bin/systemctl",
		"show", sliceName)

	outBytes, err := cmdSystemctl.CombinedOutput()
	if err != nil {
		return tasks, fmt.Errorf("Failed to run systemctl command: /usr/bin/systemctl show %s, %s, %s", sliceName, string(outBytes), err)
	}

	for _, line := range strings.Split(string(outBytes), "\n") {
		if strings.HasPrefix(line, "RequiredBy=") {
			required := strings.TrimPrefix(line, "RequiredBy=")

			for _, u := range strings.Split(required, " ") {
				t, err := StatusTask(sliceName, u, outputDir)
				if err != nil {
					errs = fmt.Sprintf("%s\n%s", errs, err)
				} else {
					tasks = append(tasks, t)
				}
			}
		}
	}

	if len(errs) > 0 {
		return tasks, fmt.Errorf("Failed to get status for all running tasks:\n%s", errs)
	}

	return tasks, nil
}

// Get status for specified task
// If outputDir is empty string
// stdout and stderr will not
// be read
func StatusTask(sliceName, unitName, outputDir string) (common.NodeTask, error) {
	task := common.NodeTask{}
	nameSplit := strings.Split(unitName, "-")

	if len(nameSplit) == 5 {
		j, err := strconv.ParseUint(nameSplit[3], 10, 64)
		if err != nil {
			return task, fmt.Errorf("Failed to parse jobid from systemd unitname: %s", unitName)
		}

		task.Task.JobID = uint(j)

		taskidSplit := strings.Split(nameSplit[4], ".")
		if len(taskidSplit) == 2 {
			t, err := strconv.Atoi(taskidSplit[0])
			if err != nil {
				return task, fmt.Errorf("Failed to convert taskid: %s to int in systemd unitname: %s", taskidSplit[0], unitName)
			}

			task.Task.Id = t
		} else {
			return task, fmt.Errorf("Failed to parse taskid from systemd unitname: %s", unitName)
		}

		task.Task.Unitname = unitName

		cmdSystemdTask := exec.Command("/usr/bin/systemctl",
			"show", unitName)

		outBytes, err := cmdSystemdTask.CombinedOutput()
		if err != nil {
			return task, fmt.Errorf("Failed to run systemctl show for task: %s, %s", unitName, err)
		}

		for _, line := range strings.Split(string(outBytes), "\n") {
			if strings.HasPrefix(line, "UID=") && line != "UID=[not set]" {
				uid, _ := strings.CutPrefix(line, "UID=")
				u, err := user.LookupId(strings.TrimSpace(uid))
				if err != nil {
					return task, fmt.Errorf("Failed to lookup user with uid: %s, %s", uid, err)
				}

				task.User = u.Username
			}

			if strings.HasPrefix(line, "GID=") && line != "GID=[not set]" {
				gid, _ := strings.CutPrefix(line, "GID=")
				g, err := user.LookupGroupId(strings.TrimSpace(gid))
				if err != nil {
					return task, fmt.Errorf("Failed to lookup group with gid: %s, %s", gid, err)
				}

				task.Group = g.Name
			}

			if strings.HasPrefix(line, "MainPID=") {
				pidStr, _ := strings.CutPrefix(line, "MainPID=")
				pid, err := strconv.Atoi(strings.TrimSpace(pidStr))
				if err == nil {
					task.Task.MainPid = pid
				}
			}

			if strings.HasPrefix(line, "ExecMainStartTimestamp=") {
				start, _ := strings.CutPrefix(line, "ExecMainStartTimestamp=")

				startSplit := strings.Split(start, " ")
				if len(startSplit) == 4 {
					task.Task.StartTime = fmt.Sprintf("%s %s", startSplit[1], startSplit[2])
				}
			}

			if strings.HasPrefix(line, "ExecMainStartTimestampMonotonic=") {
				startStr, _ := strings.CutPrefix(line, "ExecMainStartTimestampMonotonic=")
				monoStart, err := strconv.Atoi(strings.TrimSpace(startStr))
				if err == nil {
					task.Task.MonoStartTime = monoStart
				}
			}

			if strings.HasPrefix(line, "ExecMainExitTimestamp=") {
				end, _ := strings.CutPrefix(line, "ExecMainExitTimestamp=")

				endSplit := strings.Split(end, " ")
				if len(endSplit) == 4 {
					task.Task.EndTime = fmt.Sprintf("%s %s", endSplit[1], endSplit[2])
				}
			}

			if strings.HasPrefix(line, "ExecMainExitTimestampMonotonic=") {
				endStr, _ := strings.CutPrefix(line, "ExecMainExitTimestampMonotonic=")
				monoEnd, err := strconv.Atoi(strings.TrimSpace(endStr))
				if err == nil {
					task.Task.MonoEndTime = monoEnd
				}
			}

			if strings.HasPrefix(line, "ActiveState=") {
				activeState, _ := strings.CutPrefix(line, "ActiveState=")
				task.Task.ActiveState = strings.TrimSpace(activeState)
			}

			if strings.HasPrefix(line, "SubState=") {
				subState, _ := strings.CutPrefix(line, "SubState=")
				task.Task.SubState = strings.TrimSpace(subState)
			}

			if strings.HasPrefix(line, "ExecMainStatus=") {
				exitStr, _ := strings.CutPrefix(line, "ExecMainStatus=")

				if len(exitStr) > 0 {
					exitCode, err := strconv.Atoi(strings.TrimSpace(exitStr))
					if err == nil {
						task.Task.ExitCode = exitCode
					}
				}
			}

			memPeakBytes, err := os.ReadFile(fmt.Sprintf("/sys/fs/cgroup/%s/%s/memory.peak",
				sliceName, unitName))
			if err == nil {
				memPeak, err := strconv.ParseUint(strings.TrimSpace(string(memPeakBytes)), 10, 64)
				if err == nil {
					task.Task.MemoryPeak = memPeak
				}
			}

			pidsPeakBytes, err := os.ReadFile(fmt.Sprintf("/sys/fs/cgroup/%s/%s/pids.peak",
				sliceName, unitName))
			if err == nil {
				pidsPeak, err := strconv.Atoi(strings.TrimSpace(string(pidsPeakBytes)))
				if err == nil {
					task.Task.PidsPeak = pidsPeak
				}
			}
		}

		if len(outputDir) > 0 {
			stdoutPath := filepath.Join(outputDir,
				fmt.Sprintf("%s.stdout", unitName))
			stdoutBytes, err := os.ReadFile(stdoutPath)
			if err != nil {
				return task, fmt.Errorf("Failed to read stdout for jobid: %d, taskid: %d, file: %s, %s", task.Task.JobID, task.Task.Id, stdoutPath, err)
			}

			task.Task.Stdout = string(stdoutBytes)

			stderrPath := filepath.Join(outputDir,
				fmt.Sprintf("%s.stderr", unitName))
			stderrBytes, err := os.ReadFile(stderrPath)
			if err != nil {
				return task, fmt.Errorf("Failed to read stderr for jobid: %d, taskid: %d, file: %s, %s", task.Task.JobID, task.Task.Id, stderrPath, err)
			}

			task.Task.Stderr = string(stderrBytes)
		}
	} else {
		return task, fmt.Errorf("Failed to parse jobid and taskid from systemd unitname: %s", unitName)
	}

	return task, nil
}
